import os
import shutil

# === 路径配置 ===
IMAGE_SRC_DIR = "/data/suziren25/suziren/nnUNet/nnUNet_raw/abdominal_US/abdominal_US/AUS/images/train"        # 原始图像目录
MASK_ROOT_DIR = "/data/suziren25/suziren/nnUNet/nnUNet_raw/abdominal_US/abdominal_US_organs/labelsTr"         # 各器官mask的根目录
OUTPUT_ROOT_DIR = "/data/suziren25/suziren/nnUNet/nnUNet_raw/abdominal_US/abdominal_US_organs/imagesTr"       # 各器官图像输出根目录

os.makedirs(OUTPUT_ROOT_DIR, exist_ok=True)

total_copied = 0
organs_processed = 0

# 遍历所有器官文件夹
for organ in os.listdir(MASK_ROOT_DIR):
    organ_mask_dir = os.path.join(MASK_ROOT_DIR, organ)
    if not os.path.isdir(organ_mask_dir):
        continue

    organ_output_dir = os.path.join(OUTPUT_ROOT_DIR, organ)
    os.makedirs(organ_output_dir, exist_ok=True)

    copied_count = 0
    for mask_file in os.listdir(organ_mask_dir):
        if not mask_file.endswith(".png"):
            continue

        src_image_path = os.path.join(IMAGE_SRC_DIR, mask_file)
        dst_image_path = os.path.join(organ_output_dir, mask_file)

        if os.path.exists(src_image_path):
            shutil.copy(src_image_path, dst_image_path)
            copied_count += 1
            total_copied += 1

    organs_processed += 1
    print(f"✅ {organ}: 已复制 {copied_count} 张原图到 {organ_output_dir}")

print(f"\n🎯 全部完成，共处理 {organs_processed} 个器官，总计复制 {total_copied} 张图片。")
